(ns chart.jfree
  (:use [clojure.math.numeric-tower :only [ceil]]
        [clojure.java.io :only [file]])
  (:require clojure.string)
  (:import [org.jfree.chart ChartFactory ChartFrame JFreeChart ChartUtilities]
           [org.jfree.chart.axis CategoryLabelPositions]
           [org.jfree.chart.plot XYPlot PlotOrientation]
           [org.jfree.data.general DefaultPieDataset]
           [org.jfree.data.category DefaultCategoryDataset]
           [org.jfree.data.xy XYSeries XYSeriesCollection]
           [org.jfree.data.time TimeSeries TimeSeriesCollection Second]
           [org.jfree.chart.axis DateAxis SymbolAxis]
           [org.jfree.chart.renderer.xy XYBarRenderer]
           [org.jfree.data.statistics HistogramDataset]
           [org.jfree.data.gantt TaskSeriesCollection TaskSeries Task XYTaskDataset]))

;;; general

(defn show [title chart]
  (doto (ChartFrame. title chart)
    (.pack)
    (.setVisible true)))

(defn- guess-type [filename]
  (let [mapping {"PNG" :png
                 "JPEG" :jpeg
                 "JPG" :jpeg}
        ext (.toUpperCase (clojure.string/replace-first filename #"^.*\." ""))
        type (get mapping ext)]
    (when-not type
      (throw (Exception. (str "Cannot guess the image type of " filename))))
    type))

(defmacro save-png [& args]
  `(ChartUtilities/saveChartAsPNG ~@args))

(defmacro save-jpeg [& args]
  `(ChartUtilities/saveChartAsJPEG ~@args))

(def default-width 800)
(def default-height 600)

(defn save
  ([chart file-type path width height]
    (condp = file-type
      :png (save-png (file path) chart width height)
      :jpeg (save-jpeg (file path) chart width height)
      (throw (Exception. (str "Unsupported image type " file-type)))))
  ([chart path width height]
    (save chart (guess-type path) path width height))
  ([chart file-type path]
    (save chart file-type path default-width default-height))
  ([chart path]
    (save chart (guess-type path) path)))

;;; axis

(defn date-axis [label]
  (DateAxis. label))

(defn symbol-axis [label symbols]
  (SymbolAxis. label (into-array String symbols)))

;;; pie chart

(def ^:dynamic *display-chart-legend* false)
(def ^:dynamic *display-chart-tooltips* true)
(def ^:dynamic *display-chart-url* false)

(defn pie-chart [title data]
  (ChartFactory/createPieChart title data
    *display-chart-legend* *display-chart-tooltips* *display-chart-url*))

(defn pie-dataset [values]
  (let [data (DefaultPieDataset.)]
    (doseq [[k v] values]
      (.setValue data k v))
    data))

;;; line chart

(defn line-chart [title x y data]
  (ChartFactory/createLineChart title x y data PlotOrientation/VERTICAL
    *display-chart-legend* *display-chart-tooltips* *display-chart-url*))

(defn line-dataset [values]
  (let [data (DefaultCategoryDataset.)]
    (doseq [[v r c] values]
      (.addValue data v r c))
    data))

(defn line1-dataset [values]
  (line-dataset (for [[k v] values] [v "" k])))

(defn linen-dataset [values]
  (line-dataset
    (apply concat
      (for [[label values] values]
        (for [[x y] values]
          [y label x])))))

;;; Histogram

(defn histogram-chart [title x y data]
  (ChartFactory/createHistogram title x y data PlotOrientation/VERTICAL
    *display-chart-legend* *display-chart-tooltips* *display-chart-url*))

(defn histogram-dataset [bins values]
  (doto (HistogramDataset.)
    (.addSeries "" (double-array values) bins)))

;;; XY

(defn xy-chart [title x y data]
  (ChartFactory/createXYLineChart title x y data PlotOrientation/VERTICAL
    *display-chart-legend* *display-chart-tooltips* *display-chart-url*))

(defn xy-dataset [values]
  (let [data (XYSeriesCollection.)]
    (doseq [[label values] values]
      (let [xy (XYSeries. label)]
        (doseq [[x y] values]
          (.add xy (double x) (double y)))
        (.addSeries data xy)))
    data))

(defn xy-plot [dataset domainAxis rangeAxis renderer]
  (XYPlot. dataset domainAxis rangeAxis renderer))

;;; Time

(defn time-chart [title x y data]
  (ChartFactory/createTimeSeriesChart title x y data
    *display-chart-legend* *display-chart-tooltips* *display-chart-url*))

(defn time-dataset [values]
  (let [series (TimeSeries. "")]
    (doseq [[time value] values]
      (.addOrUpdate series (Second. time) value)) ; overwrite duplicated records, though I don't understand why they exist...
    (TimeSeriesCollection. series)))

;;; interval

(defn interval-xy-dataset [values]
  (let [data (TaskSeriesCollection.)]
    (doseq [[label values] values]
      (let [ts (TaskSeries. label)]
        (doseq [[start end] values]
          (letfn [(to-date [x] (if (integer? x) (java.util.Date. x) x))]
            (.add ts (Task. "" (to-date start) (to-date end)))))
        (.add data ts)))
    (let [dataset (XYTaskDataset. data)]
      (.setTransposed dataset true)
      ;(.setSeriesWidth dataset 0.5)
      dataset)))

(defn interval-plot [x-label y-label values]
  (xy-plot (interval-xy-dataset values)
    (date-axis x-label)
    (symbol-axis y-label (map first values))
    (doto (XYBarRenderer.) (.setUseYInterval true))))

(defn interval-chart [title x y data]
  (JFreeChart. title JFreeChart/DEFAULT_TITLE_FONT (interval-plot x y data) *display-chart-legend*))


;;; chart && dataset

(defn chart [chart-type & args]
  (apply (condp = chart-type
    :pie pie-chart
    :line line-chart
    :xy xy-chart
    :time time-chart
    :interval interval-chart
    :histogram histogram-chart
    ) args))

(defn dataset [chart-type values]
  (let [fname (symbol (str (name chart-type) "-dataset"))]
    (if-let [v (ns-resolve 'chart.jfree fname)]
      ((var-get v) values)
      (throw (Exception. (format "%s for chart type %s is not found" fname chart-type))))))

;;; look and feel

(defn- to-font-style [style]
  (condp = style
    :plain java.awt.Font/PLAIN
    :bold java.awt.Font/BOLD
    style))

(def default-chart-font-size 12)

(defn make-font
  ([name style size]
    (java.awt.Font. name (to-font-style style) size))
  ([name-or-seq]
    (cond
      (string? name-or-seq) (make-font name-or-seq :plain default-chart-font-size)
      (sequential? name-or-seq) (apply make-font name-or-seq)
      (instance? java.awt.Font name-or-seq) name-or-seq
      :else (throw (Exception. (str "Don't know how to create font from " name-or-seq)))))
  ([name style-or-size]
    (if (number? style-or-size)
      (make-font name :plain style-or-size)
      (make-font name style-or-size default-chart-font-size))))

(defn use-font
  ([chart font0 font1 font2 font3 font4]
    (let [plot (.getPlot chart)
          domain-axis (.getDomainAxis plot)
          range-axis  (.getRangeAxis  plot)]
      (when font0 (.. chart getLegend (setItemFont (make-font font0))))
      (when font1 (.setTickLabelFont domain-axis font1))
      (when font2 (.setLabelFont     domain-axis font2))
      (when font3 (.setTickLabelFont range-axis  font3))
      (when font4 (.setLabelFont     range-axis  font4)))
    chart)
  ([chart font0 font]
    (use-font chart font0 font font font font))
  ([chart font]
    (use-font chart nil font font font font)))

(defn use-chinese [chart]
  (use-font chart (make-font "\u5b8b\u4f53"))
  chart)

(defn set-x-date-format
  ([chart format]
    (let [plot (.getPlot chart)
          domain-axis (.getDomainAxis plot)]
      (.setDateFormatOverride domain-axis (java.text.SimpleDateFormat. format))
      chart))
  ([chart]
    (set-x-date-format chart "MM-dd")))

(defn set-y-from-zero [chart]
  (let [plot (.getPlot chart)
        range-axis (.getRangeAxis plot)]
    (.setAutoRangeIncludesZero range-axis true)
    chart))

(defn set-y-range [chart from to]
  (let [plot (.getPlot chart)
        range-axis (.getRangeAxis plot)]
    (.setRange range-axis from to)
    chart))

(defn- to-color [color]
  (condp = color
    :white java.awt.Color/WHITE
    :light_gray java.awt.Color/LIGHT_GRAY
    :gray java.awt.Color/GRAY
    :dark_gray java.awt.Color/DARK_GRAY
    :black java.awt.Color/BLACK
    :red java.awt.Color/RED
    :pink java.awt.Color/PINK
    :orange java.awt.Color/ORANGE
    :yellow java.awt.Color/YELLOW
    :green java.awt.Color/GREEN
    :magenta java.awt.Color/MAGENTA
    :cyan java.awt.Color/CYAN
    :blue java.awt.Color/BLUE
    color))

(defn use-color
  ([chart index color]
    (.. chart getPlot getRenderer (setSeriesPaint index (to-color color)))
    chart)
  ([chart color]
    (use-color chart 0 color)))

;;; Utilities

(def ^:dynamic *fool-max* 10)

(defn fool
  ([values max]
    (let [n (count values)
          step (ceil (/ n max))
          x-values (map #(nth values %) (range 0 n step))
          values (map #(nth x-values (quot % step)) (range n))]
      values))
  ([values]
    (fool values *fool-max*)))

(defn fool1
  ([values max]
    (map cons (fool (map first values) max) (map rest values)))
  ([values]
    (fool1 values *fool-max*)))

(defn fooln
  ([values max]
    (for [[label values] values] [label (fool1 values max)]))
  ([values]
    (fooln values *fool-max*)))

;(defn rotate-x [chart]
;  (.. chart getCategoryPlot getDomainAxis (setCategoryLabelPositions (CategoryLabelPositions/createUpRotationLabelPositions (/ Math/PI 2))))
;  chart)

(defn rotate-x [chart]
  (.. chart getPlot getDomainAxis (setVerticalTickLabels true)); FIXME: not working in examples.jfree-demo
  chart)

;;; test

